import time
from PyQt5.QtCore import Qt, QThread, QMutex, QWaitCondition
import cv2
from d435 import D435
from ultralytics import YOLO
from arm import Rokae
from util_3d import *
from  grab import grab, yolomask_to_pointcloud
from pv_agv import AGV
from lift import Lift
from PyQt5.QtCore import QObject, pyqtSignal
from log import logger


class Task(QThread):
    info_signal = pyqtSignal(str)
    def __init__(self, taskConfig, arm:Rokae, camera:D435, model:YOLO, lift:Lift, controller):
        super(Task, self).__init__()
        # 设置工作状态与初始num数值
        self.working = True
        self.num = 0
        self.arm = arm
        self.camera = camera
        self.model = model
        self.lift = lift
        self.taskConfig = taskConfig
        self.restoplt_image = None
        self.mutex = QMutex()
        self.condition = QWaitCondition()
        self.is_paused = False
        self.controller = controller

        self.arm_init_pose = np.array(
            [0.23553320870561686, 0.13599433438178785, 0.33854353132547865, -3.1411474368438133,
             -0.00030757519895313585,
             2.34705996088468])

        if taskConfig is not None:
            self.agv_station_list = taskConfig["agv_station_list"]
            self.agv_user_name = taskConfig["user_name"]
            self.agv_passwd = taskConfig["passwd"]
            self.agv = AGV(taskConfig['agv_ip'], taskConfig['user_name'], taskConfig['passwd'])
            self.agv_map =  taskConfig['agv_map_name']
            self.agv.load_map(self.agv_map)
            self.lift_offset_list = taskConfig["lift_offset"]
            self.offset_y_list = taskConfig["offset_y_list"]
            self.offset_z_list = taskConfig["offset_z_list"]
            self.arm_obs_pose_list = []
            for offset_y in self.offset_y_list:
                for offset_z in self.offset_z_list:
                    arm_obs_pose = self.arm_init_pose.copy()
                    arm_obs_pose[1] =   self.arm_init_pose[1] + offset_y
                    arm_obs_pose[2] = self.arm_init_pose[2] + offset_z
                    self.arm_obs_pose_list.append(arm_obs_pose)

        self.default_pose = [0.2532188848287554, 0.09030179288176225, 0.3384137821651567, -3.1409692078286757, -0.00048430784375869286, 2.3215618536316596]
        self.place_pose = [0.10680710828535048, -0.20382760394377142, 0.3384161440699347, -3.1409714661824824, -0.00048260724761573523, 0.816773405620627]
        self.place_pose2 = [0.10680710828535048, -0.20382760394377142, 0.2384161440699347, -3.1409714661824824, -0.00048260724761573523, 0.816773405620627]
        self.arm_init_pose = np.array(self.default_pose)
        self.station_index = -1
        self.task_state = "not_started"
        self.info_signal.connect(self.controller.onInfo)
        # Y 方向这 Z方向移动

    def __del__(self):
        # 线程状态改变与线程终止
        self.working = False
        self.wait()

    def setTaskParam(self, taskConfig):
        pass
    def print_info(self, info):
        # 发射信号并传递参数
        # logger.info(info)
        self.info_signal.emit(info)

    def obs(self, dist_max, dist_mini):
        object_list = []
        self.camera.update()
        color_image, depth_image = self.camera.get_images()
        # model output
        results = self.model(color_image,conf = 0.8)
        annotation_frame = results[0].plot()
        cv2.imshow("obs", annotation_frame)
        cv2.waitKey(100)
        self.result_image = annotation_frame
        self.num +=1
        if results[0].masks is not None:
            name_dict = results[0].names
            for i in range(len(results[0].masks)):
                mask_np = results[0].masks[i].data.cpu().numpy()
                class_index = results[0].boxes.cls[i].data.cpu().numpy()
                if class_index == 0:
                    pcd = yolomask_to_pointcloud(color_image, depth_image, mask_np)
                    center = pcd.get_center()
                    cam_aim_point = np.zeros(4)
                    cam_aim_point[:3] = center
                    cam_aim_point[3] = 1
                    if dist_max > cam_aim_point[2] > dist_mini:
                        object_list.append(cam_aim_point)
                        print(center)
        return  object_list


    def oneStep(self):

        self.num = self.num + 1
        end_t_cam = halcon_pose_to_T(self.camera.out_param_halcon_type)
        # move to obs pose
        self.arm.setPose2(lambda st: True, 50, 50, 50, self.init_pos.tolist())
        # update image
        self.camera.update()
        color_image, depth_image = self.camera.get_images()
        # model output
        results = self.model(color_image)
        annotation_frame = results[0].plot()
        self.result_image = annotation_frame
        if results[0].masks is not None:
            name_dict = results[0].names
            for i in range(len(results[0].masks)):
                mask_np = results[0].masks[i].data.cpu().numpy()
                class_index = results[0].boxes.cls[i].data.cpu().numpy()
                if class_index == 47:
                    pcd = yolomask_to_pointcloud(color_image, depth_image, mask_np)
                    center = pcd.get_center()
                    cam_aim_point = np.zeros(4)
                    cam_aim_point[:3] = center
                    cam_aim_point[3] = 1
                    current_pose = self.arm.getPose()
                    base_t_end = rokae_pose_to_T(current_pose)
                    cam_t_tool = transl(0.04, 0.08, 0.10)
                    base_t_aim_end = grab(base_t_end, end_t_cam, cam_t_tool, cam_aim_point)
                    base_t_aim_end_rokae = T_to_rokae_pose(base_t_aim_end)
                    base_t_aim_end_front_rokae =  base_t_aim_end_rokae.copy()
                    base_t_aim_end_front_rokae[0] = base_t_aim_end_front_rokae[0] - 0.15

                    self.arm.setPose2(lambda st: True, 50, 50, 50, base_t_aim_end_front_rokae)
                    self.arm.setEndOpen()
                    self.arm.setPose2(lambda st: True, 50, 50, 50, base_t_aim_end_rokae)
                    self.arm.setEndClose()
                    self.arm.setPose2(lambda st: True, 50, 50, 50, base_t_aim_end_front_rokae)
                    self.arm.setEndOpen()
                    self.arm.setPose2(lambda st: True, 50, 50, 50, self.init_pos.tolist())
        pass
    # cam 下的目标点
    def oneStepArmGrab(self, cam_aim_point_list):
        self.num = self.num + 1
        end_t_cam = halcon_pose_to_T(self.camera.out_param_halcon_type)
        # move to obs pose
        # self.arm.setPose2(lambda st: True, 50, 50, 50, grab_pose.tolist())
        # update image
        # model output
        if len(cam_aim_point_list) !=0:
            cam_aim_point = cam_aim_point_list[0]
            current_pose = self.arm.getPose()
            base_t_end = rokae_pose_to_T(current_pose)
            cam_t_tool = transl(0.04, 0.08, 0.10)
            base_t_aim_end = grab(base_t_end, end_t_cam, cam_t_tool, cam_aim_point)
            base_t_aim_end_rokae = T_to_rokae_pose(base_t_aim_end)
            base_t_aim_end_rokae[2] = base_t_aim_end_rokae[2] + 0.03
            base_t_aim_end_front_rokae = base_t_aim_end_rokae.copy()
            base_t_aim_end_front_rokae[0] = base_t_aim_end_front_rokae[0] - 0.15

            self.arm.setPose2(lambda st: True, 100, 100, 100, base_t_aim_end_front_rokae)
            self.arm.setEndOpen()
            self.arm.setPose2(lambda st: True, 100, 100, 100, base_t_aim_end_rokae)
            self.arm.setEndClose()
            self.arm.setPose2(lambda st: True, 100, 100, 100, base_t_aim_end_front_rokae)
            # self.arm.setEndOpen()
            self.arm.setPose2(lambda st: True, 100, 100, 100, current_pose)
            self.arm.setPose2(lambda st: True, 100, 100, 100, self.default_pose)
            self.arm.setPose2(lambda st: True, 100, 100, 100, self.place_pose)
            self.arm.setPose2(lambda st: True, 100, 100, 100, self.place_pose2)
            self.arm.setEndOpen()
            # self.lift.set_pos(0.0)
            time.sleep(1)
            self.arm.setPose2(lambda st: True, 100, 100, 100, self.place_pose)
            self.arm.setPose2(lambda st: True, 100, 100, 100, self.default_pose)
            # self.lift.set_pos(0.0)
            time.sleep(1)
            pass

    def setObsPoseList(self, poseList):
        pass

    def oneStepAll(self):
        if self.station_index == -1 and self.task_state == "not_started":
            # self.arm.setPose2(lambda st: True, 50, 50, 50, self.arm_init_pose.tolist())
            self.arm.setPose2(lambda st: True, 50, 50, 50, self.default_pose)
            self.station_index = 0
            self.agv.load_map(self.agv_map)
            self.agv.move_to_station_list([self.agv_station_list[self.station_index]])
            self.print_info("移动到目标位置"+str(self.station_index) + self.agv_station_list[self.station_index])
            self.task_state = "running"
        else:
            if self.agv.is_idle():
                self.print_info("到达目标地点")
                if self.agv_station_list[self.station_index].startswith("c_"):
                    for lift_offset in self.lift_offset_list:
                        self.print_info("升降台移动到观察位置"+str(lift_offset))
                        self.lift.set_pos_block(lift_offset)
                        for arm_pose in self.arm_obs_pose_list :
                            self.arm.setPose2(lambda st: True, 50, 50, 50, arm_pose.tolist())
                            self.print_info("机械臂移动到观察位置")
                            while True:
                                obs_res = self.obs(0.7,0.2)
                                if obs_res is not None and len(obs_res) > 0:
                                    self.oneStepArmGrab(obs_res)
                                    if not self.working:
                                        return
                                else:
                                    self.print_info("未发现可采摘物体")
                                    break
                    self.print_info("升降台回收")
                    self.lift.set_pos_block(0.0)
                    # 移动到下一个站点
                    self.station_index +=1
                    if  self.station_index >= len(self.agv_station_list):
                        self.task_state = "end"
                        self.print_info("任务结束")
                        self.working = False
                    else:
                        self.agv.move_to_station_list([self.agv_station_list[self.station_index]])
                        self.print_info("移动到目标位置" + self.agv_station_list[self.station_index])
                        time.sleep(2)
            else: # agv is moving
                self.print_info( "agv 正在移动")

    def oneStepTest(self):
        self.arm.setPose2(lambda st: True, 50, 50, 50, self.default_pose)
        self.lift.set_pos_block(100)
        obs_res = self.obs(0.5, 0.2)
        if obs_res is not None and len(obs_res) > 0:
            self.oneStepArmGrab(obs_res)

    def run(self):
        while self.working :
            self.mutex.lock()
            if self.is_paused:
                self.condition.wait(self.mutex)
            if self.taskConfig is not None:
                self.mutex.unlock()
                self.oneStepAll()
            else:
                self.mutex.unlock()
                self.oneStepAll()
    def pause(self):
        self.mutex.lock()
        self.is_paused = True
        self.mutex.unlock()
        self.print_info("任务暂停")

    def resume(self):
        self.mutex.lock()
        self.is_paused = False
        self.condition.wakeAll()
        self.mutex.unlock()
        self.print_info("任务继续")
    def stop(self):
        self.working = False
        self.print_info("任务结束")
